#include "fe1dParabolicTask.h"
#include "feMatrix.h"
#include "feMeshFunction.h"
#include "feLuSolver.h"
#include "feObserver.h"
#include "fePolylineGenerator.h"
#include "feElement.h"
#include "feBordCond.h"
#include "feVector.h"
#include "feMaterial.h"
#include "../Timer.h"

#include <math.h>

fe1dParabolicTask::fe1dParabolicTask()
{
	X0 = 0;
	X1 = 0;
	X2 = 0;
	Solver = new feLuSolver();
	Right = 0;
}


//----------------------------------------------------------------------------------------------------------------------------------------------
bool fe1dParabolicTask::Init(feMesh *ioMesh, feBordCond **iConds, int iCondsCount, feVector *iX0, double iFirstTimeStep, 
							 int iLayersTime, double iStepCoeffTime)
{
	FirstTimeStep = iFirstTimeStep;
	CoeffTime = iStepCoeffTime;
	LayersTime = iLayersTime;
	CurLayer = 0;
	Mesh = ioMesh;
	Bcs = iConds;
	NumBc = iCondsCount;

	
	//allocate
	if(!X0) X0 = new feVector(Mesh->NodesCount());
	if(!X1) X1 = new feVector(Mesh->NodesCount());
	if(!X2) X2 = new feVector(Mesh->NodesCount());
	
	F.Allocate(Mesh->NodesCount());
	M.Allocate(Mesh->NodesCount());
	Stiffness.Allocate(Mesh->NodesCount()); 
	WeightWithDuDt.Allocate(Mesh->NodesCount()); 
	WeightWithGamma.Allocate(Mesh->NodesCount());
		

	//copy x0
	if(iX0)
		feMvOps::CopyVV(iX0, X0);
	
	Right = (feMeshFunction*)0;

	return true;
}


//---------------------------------------------------------------------------------------------------------------------------------------------
bool fe1dParabolicTask::Init(feMesh *ioMesh, feBordCond **iConds, int iCondsCount, feVector *iX0, 
							feMeshFunction *iRight, int iLayersPass, int iNumL, double iScaleCoeff)
{
	LayersTime = iNumL;
	LayersPass = iLayersPass;
	CoeffTime = iScaleCoeff;

	CurLayer = 0;
	Mesh = ioMesh;
	Bcs = iConds;
	NumBc = iCondsCount;
	Right = iRight;


	//allocate
	if(!X0) X0 = new feVector(Mesh->NodesCount());
	if(!X1) X1 = new feVector(Mesh->NodesCount());
	if(!X2) X2 = new feVector(Mesh->NodesCount());

	F.Allocate(Mesh->NodesCount());
	M.Allocate(Mesh->NodesCount());
	Stiffness.Allocate(Mesh->NodesCount()); 
	WeightWithDuDt.Allocate(Mesh->NodesCount()); 
	WeightWithGamma.Allocate(Mesh->NodesCount());


	//copy x0
	if(iX0)
		feMvOps::CopyVV(iX0, X0);

	return true;
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::RunSolver()
{
	double eps = 1;
	int iters = 0;
	
	eps = Solver->GetEps();
	for(; iters<SOLVER_MAX_ITER && eps > SOLVER_EPS; iters++)
	{
		Solver->Iterate();
		eps = Solver->GetEps();
		if(iters % 10 == 0 && iters > 0)
		{
			char str[256];
			sprintf(str, "Iter  = %d, eps = %e\n", iters, eps);
			feObserverManager::GetInstance().Notify(str);
		}
	}
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::ProcessBorderConds()
{
	//process all border conds
	for(int i=0; i<NumBc; i++)
		if(Bcs[i]) Bcs[i]->AddToSlae(&M, &F, CurLayer);	
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::GatherStiffness()
{
	//zero stiffness
	Stiffness.Zero();


	feDenseMatrix *s = NULL;
	feReal cobuf[3];
	feVector coeffs(3, cobuf);
	int remap[100];


	//iterate over elements
	feMesh::ElementIterator it = Mesh->GetIterator();
	for(it.First(); !it.IsDone(); it.Next())
	{
		//get element
		feElement *el = it.CurrentItem();

		//get material coeffs:  divgrad, weight and dudt 
		el->GetMaterial()->GetCoeffs(&coeffs);

		//get local stiffness
		el->GetStiffness(&s);

		//get remap
		el->GetIndeces(remap);

		//add local to global
		feMvOps::AddLocalToGlobalM(&Stiffness, s, remap, coeffs.Data[0]);
	}

	if(s)
		delete s;
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::GatherWeight()
{
	//zero weight matrices
	WeightWithGamma.Zero();
	WeightWithDuDt.Zero();
	

	feDenseMatrix *w = NULL;
	feReal cobuf[3];
	feVector coeffs(3, cobuf);
	int remap[100];

	//iterate over elements
	feMesh::ElementIterator it = Mesh->GetIterator();
	for(it.First(); !it.IsDone(); it.Next())
	{
		//get element
		feElement *el = it.CurrentItem();

		//get material coeffs:  divgrad, weight and dudt 
		el->GetMaterial()->GetCoeffs(&coeffs);

		//get locals
		el->GetWeight(&w);


		//get remap
		el->GetIndeces(remap);


		//add local to global
		feMvOps::AddLocalToGlobalM(&WeightWithGamma, w, remap, coeffs.Data[1]);		
		feMvOps::AddLocalToGlobalM(&WeightWithDuDt, w, remap, coeffs.Data[2]);
	}

	if(w) delete w;
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::GatherRight()
{
	//gather right
	feReal fbuf[100];
	feVector f(3, fbuf);
	int remap[100];

	feMesh::ElementIterator it = Mesh->GetIterator();
	for(it.First(); !it.IsDone(); it.Next())
	{
		//get element
		feElement *el = it.CurrentItem();

		//get locals
		el->GetRight(&f, CurLayer);

		//get remap
		el->GetIndeces(remap);

		//add local to global
		feMvOps::AddLocalToGlobalV(&F, &f, remap);
	}
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::ZeroData()
{
	//zero f
	ZeroMemory(&F.Data[0], F.GetSize()*sizeof(feReal));

	//zero M
	M.Zero();
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::SolveNoLayers()
{
	//zero data
	ZeroData();

	//add stiffness
	M.AddM(&Stiffness);

	//add weight
	M.AddM(&WeightWithGamma);

	//add right
	GatherRight();


	//iterate over border conditions
	ProcessBorderConds();


	//solve
	Solver->BindData(&M, &F, X0);
	RunSolver();
	Solver->GetX(X1);
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::Solve2Layer(double iDt)
{
	//zero data
	ZeroData();
	
	//add stiffness
	M.AddM(&Stiffness);
	
	
	double coeff = 1/iDt;

	//add weight
	M.AddM(&WeightWithGamma);
	M.AddM(&WeightWithDuDt, coeff);

		
	//add right		
	GatherRight();


	//mult W*X0
	feVector mx(Mesh->NodesCount());
	feMvOps::MultMV(&WeightWithDuDt, X0, &mx);
	
	//add to global F
	feMvOps::MultAddVV(&F, &mx, &F, coeff);
	mx.Free();


	//iterate over border conditions
	ProcessBorderConds();


	//solve
	Solver->BindData(&M, &F, X0);
	RunSolver();
	Solver->GetX(X1);
}


//---------------------------------------------------------------------------------------------------------------------------------------------
void fe1dParabolicTask::Solve3Layer(double iDt1, double iDt2)
{
	//zero data
	ZeroData();


	//get dts
	double dt = iDt1 + iDt2;
	double dt0 = iDt2;
	double dt1 = iDt1;


	//add stiffness
	M.AddM(&Stiffness);


	//add weight
	M.AddM(&WeightWithGamma);
	M.AddM(&WeightWithDuDt, (dt+dt0)/(dt*dt0));


	//add right		
	GatherRight();


	//mult W*X0
	feVector mx(Mesh->NodesCount());
	feMvOps::MultMV(&WeightWithDuDt, X0, &mx);

	//add to global F
	feMvOps::MultAddVV(&F, &mx, &F, -dt0/(dt*dt1));


	//mult W*X1
	feMvOps::MultMV(&WeightWithDuDt, X1, &mx);


	//add to global F
	feMvOps::MultAddVV(&F, &mx, &F, dt/(dt1*dt0));
	mx.Free();


	//iterate over border conditions
	ProcessBorderConds();
	

	//solve
	Solver->BindData(&M, &F, X1);
	RunSolver();
	Solver->GetX(X2);
	
}


//---------------------------------------------------------------------------------------------------------------------------------------------
bool fe1dParabolicTask::SolveStationary(feVector *oSolvation, double *oFdotX)
{
	//check mesh
	if(!Mesh) return false;

	//gather global matrices once
	GatherStiffness();
	GatherWeight();


	//solve
	SolveNoLayers();


	//copy
	if(oSolvation)
		feMvOps::CopyVV(X1, oSolvation);


	if(oFdotX)
		(*oFdotX) = feMvOps::DotVV(X1, &F);

	//done
	return true;
}


//---------------------------------------------------------------------------------------------------------------------------------------------
bool fe1dParabolicTask::SolveNonStationary(feMeshFunction *oFunc)
{
	//gather global matrices once
	GatherStiffness();
	GatherWeight();


	//check mesh
	if(oFunc->GetMesh() != Mesh)
		return false;
	

	//check mesh
	if(!Mesh) return false;

	
	//get start time
	double time = FirstTimeStep;

	if(Right != 0)
	{
		CurLayer = LayersPass;
		time = Right->GetTime(LayersPass) / CoeffTime;
	}
	else 
		CurLayer = 1;

	
	//run 2 layer once
	Solve2Layer(time);


	//add 2 layers to mesh
	oFunc->AddLayer(0, X0);
	oFunc->AddLayer(time, X1);


	//run 3 layer for all other
	double dt1 = time;
	double dt2 = time;


	for(int i=2; i<LayersTime; i++)
	{
		//update dt2
		if((feMeshFunction*)Right != 0)
		{
			dt2 = Right->GetTime(i*LayersPass) / CoeffTime - time;
			CurLayer = i*LayersPass;
		}
		else
		{
			dt2 *= CoeffTime;
			CurLayer++;
		}

		time += dt2;
				

		//solve
		Solve3Layer(dt1, dt2);

		//add layer to mesh
		oFunc->AddLayer(time, X2);

		//shift
		feVector *b = X0;
		X0 = X1;
		X1 = X2;
		X2 = b;
			
		//update dt1
		dt1 = dt2;
	}


	//done
	return true;
}


//---------------------------------------------------------------------------------------------------------------------------------------------
fe1dParabolicTask::~fe1dParabolicTask()
{
	if(X0) delete X0;
	if(X1) delete X1;
	if(X2) delete X2;
	if(Solver)	delete Solver; 
}